import numpy as np
import matplotlib.pyplot as plt
import pylab

coordinates = np.array([[565.0, 575.0], [25.0, 185.0], [345.0, 750.0], [945.0, 685.0], [845.0, 655.0],
                        [880.0, 660.0], [25.0, 230.0], [525.0, 1000.0], [580.0, 1175.0], [650.0, 1130.0],
                        [1605.0, 620.0], [1220.0, 580.0], [1465.0, 200.0], [1530.0, 5.0], [845.0, 680.0],
                        [725.0, 370.0], [145.0, 665.0], [415.0, 635.0], [510.0, 875.0], [560.0, 365.0],
                        [300.0, 465.0], [520.0, 585.0], [480.0, 415.0], [835.0, 625.0], [975.0, 580.0],
                        [1215.0, 245.0], [1320.0, 315.0], [1250.0, 400.0], [660.0, 180.0], [410.0, 250.0],
                        [420.0, 555.0], [575.0, 665.0], [1150.0, 1160.0], [700.0, 580.0], [685.0, 595.0],
                        [685.0, 610.0], [770.0, 610.0], [795.0, 645.0], [720.0, 635.0], [760.0, 650.0],
                        [475.0, 960.0], [95.0, 260.0], [875.0, 920.0], [700.0, 500.0], [555.0, 815.0],
                        [830.0, 485.0], [1170.0, 65.0], [830.0, 610.0], [605.0, 625.0], [595.0, 360.0],
                        [1340.0, 725.0], [1740.0, 245.0]])  # 城市坐标


def getdistmat(coordinates):
    num = coordinates.shape[0]
    distmat = np.zeros((52, 52))  # 打印一个52行52列的二维数组
    for i in range(num):
        for j in range(i, num):
            distmat[i][j] = distmat[j][i] = np.linalg.norm(coordinates[i] - coordinates[j])
    return distmat  # distmat[i][j] = distmat[j][i]表示城市i和j距离


distmat = getdistmat(coordinates)  # 二维数组，表示城市之间的距离
numant = 40  # 蚂蚁个数
numcity = coordinates.shape[0]  # 城市个数   shape[0]输出二维数组【numpy】的行数。
alpha = 1  # 信息素重要程度因子
beta = 5  # 启发函数重要程度因子
rho = 0.1  # 信息素的挥发速度
Q = 1  # 信息素强度
iter = 0
itermax = 250
etatable = 1.0 / (distmat + np.diag([1e10] * numcity))  # 启发函数矩阵，表示蚂蚁从城市i转移到矩阵j的期望程度
pheromonetable = np.ones((numcity, numcity))  # 信息素矩阵        52行52列的二维数组。元素都是1
pathtable = np.zeros((numant, numcity)).astype(int)  # 路径记录表    40行52列的二维数组，目前都为0
# distmat = getdistmat(coordinates)  # 城市的距离矩阵
lengthaver = np.zeros(itermax)  # 各代(各个循环)路径的平均长度          # 一维数组，itermax个0
lengthbest = np.zeros(itermax)  # 各代及其之前遇到的最佳路径长度    # 一维数组，itermax个0
pathbest = np.zeros((itermax, numcity))  # 各代及其之前遇到的最佳路径长度         # itermax(250)行numcity(52)列的二维矩阵

while iter < itermax:  # 循环0-250次
    # 1、随机产生各个蚂蚁的起点城市
    if numant <= numcity:  # 1.1、城市数比蚂蚁数多
        #  permutation()：随机排列序列，生成一个0-52的一维数组,[:numant]去除前numant个分配给pathtable中的第0列
        pathtable[:, 0] = np.random.permutation(range(0, numcity))[:numant]
    else:  # 1.2、蚂蚁数比城市数多，需要补足
        pathtable[:numcity, 0] = np.random.permutation(range(0, numcity))[:]  # 先将城市全部取出，剩下的蚂蚁随机分配
        # 剩下numant - numcity个蚂蚁，将分配给他们的城市数分配在他们的第0列
        pathtable[numcity:, 0] = np.random.permutation(range(0, numcity))[:numant - numcity]
    length = np.zeros(numant)  # 计算各个蚂蚁的路径距离，还未走,都是0.
    # 2、各个蚂蚁进行路径选择
    for i in range(numant):
        visiting = pathtable[i, 0]  # 当前所在的城市，将[i,0]位置的值赋值给visiting
        unvisited = set(range(numcity))  # 未访问的城市,以集合的形式存储{}
        unvisited.remove(visiting)  # 删除元素；利用集合的remove方法删除存储的数据内容
        # 3、循环numcity-1次，访问剩余的numcity-1个城市，循环为一只蚂蚁访问所有城市的过程
        for j in range(1, numcity):
            # 每次用轮盘法选择下一个要访问的城市
            listunvisited = list(unvisited)      # 列表,即为一维数组。
            probtrans = np.zeros(len(listunvisited))   # 创建一个【未被访问城市个数】一维数组,s元素全为0
            for k in range(len(listunvisited)):
                probtrans[k] = np.power(pheromonetable[visiting][listunvisited[k]], alpha) \
                               * np.power(etatable[visiting][listunvisited[k]], beta)
            cumsumprobtrans = (probtrans / sum(probtrans)).cumsum()    # cumsum()将probtrans数组中的值进行累加求和
            cumsumprobtrans -= np.random.rand()
            k = listunvisited[(np.where(cumsumprobtrans > 0)[0])[0]]  # python3中原代码运行bug，类型问题；鉴于此特找到其他方法
            # 通过where（）方法寻找矩阵大于0的元素的索引并返回ndarray类型，然后接着载使用[0]提取其中的元素，用作listunvisited列表中
            # 元素的提取（也就是下一轮选的城市）
            pathtable[i, j] = k  # 添加到路径表中（也就是蚂蚁走过的路径)，第i号蚂蚁第j步走的城市。
            unvisited.remove(k)  # 然后在为访问城市set中remove（）删除掉该城市
            length[i] += distmat[visiting][k]    # 当前所在城市visiting到城市k的距离，length[]用于存储各个蚂蚁走的路径长度，一维数组。
            visiting = k
        length[i] += distmat[visiting][pathtable[i, 0]]  # 蚂蚁的路径距离包括最后一个城市和第一个城市的距离  TSP问题，访问所有城市后又回到了原点
        # 包含所有蚂蚁的一个迭代结束后，统计本次迭代的若干统计参数
    # 4、从当前循环中找出最小路径和最小路径长度并保存
    lengthaver[iter] = length.mean()   # 各代(各个循环)路径的平均长度
    if iter == 0:
        lengthbest[iter] = length.min()                      # length[]记录的是当前循环，每个蚂蚁走出的路径长度
        pathbest[iter] = pathtable[length.argmin()].copy()   # length.argmin()显示length[]中最小值的下标
    else:
        if length.min() > lengthbest[iter - 1]:              # 当前循环中的最小路径不是最优值
            lengthbest[iter] = lengthbest[iter - 1]
            pathbest[iter] = pathbest[iter - 1].copy()
        else:
            lengthbest[iter] = length.min()
            pathbest[iter] = pathtable[length.argmin()].copy()
    # 5、更新信息素
    changepheromonetable = np.zeros((numcity, numcity))
    for i in range(numant):
        for j in range(numcity - 1):
            # 更新蚂蚁i从城市j到j+1路径的信息素
            changepheromonetable[pathtable[i, j]][pathtable[i, j + 1]] += Q / distmat[pathtable[i, j]][
                pathtable[i, j + 1]]  # 计算信息素增量
        # 最后一个城市到第一个城市途径的信息素增量
        changepheromonetable[pathtable[i, j + 1]][pathtable[i, 0]] += Q / distmat[pathtable[i, j + 1]][pathtable[i, 0]]
    # 带入计算信息素公式，更新完毕，本次循环结束
    pheromonetable = (1 - rho) * pheromonetable + changepheromonetable  # 计算信息素公式
    iter += 1  # 迭代次数指示器+1
    print("iter:", iter)

# 做出平均路径长度和最优路径长度
fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(12, 10))
axes[0].plot(lengthaver, 'k', marker=u'')
axes[0].set_title('Average Length')
axes[0].set_xlabel(u'iteration')

axes[1].plot(lengthbest, 'k', marker=u'')
axes[1].set_title('Best Length')
axes[1].set_xlabel(u'iteration')
fig.savefig('average_best.png', dpi=500, bbox_inches='tight')
plt.show()

# 作出找到的最优路径图
bestpath = pathbest[-1]   # 每次更新后，pathbest[-1]的最后一行，就是最优路径了
plt.plot(coordinates[:, 0], coordinates[:, 1], 'r.', marker=u'$\cdot$')
plt.xlim([-100, 2000])
plt.ylim([-100, 1500])

for i in range(numcity - 1):
    m = int(bestpath[i])
    n = int(bestpath[i + 1])
    plt.plot([coordinates[m][0], coordinates[n][0]], [coordinates[m][1], coordinates[n][1]], 'k')  # k为黑色
plt.plot([coordinates[int(bestpath[0])][0], coordinates[int(n)][0]],
         [coordinates[int(bestpath[0])][1], coordinates[int(n)][1]], 'b')  # 最后一个城市与第一个城市连线。 b为蓝色
ax = plt.gca()
ax.set_title("Best Path")
ax.set_xlabel('X axis')
ax.set_ylabel('Y_axis')

plt.savefig('best path.png', dpi=500, bbox_inches='tight')
plt.show()